package com.raos.example.chat;

import dev.langchain4j.community.model.zhipu.ZhipuAiStreamingChatModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.time.Duration;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

/**
 * 聊天流式响应
 *
 * @author raos
 * @date 2025-07-06
 */
@Slf4j
@RestController
@RequestMapping("/api/stream")
public class ChatStreamAPI {

    @Resource(name = "qwenStreamingChatModel")
    private StreamingChatLanguageModel streamingChatModel;
    @Resource(name = "ollamaStreamingChatModel")
    private StreamingChatLanguageModel streamingChatModel2;
    @Resource(name = "deepseekStreamingChatModel")
    private StreamingChatLanguageModel streamingChatModel3;

    @Value("${langchain4j.community.zhipu.streaming-chat-model.api-key}")
    private String zhipuApiKey;
    @Value("${langchain4j.community.zhipu.streaming-chat-model.model-name}")
    private String zhipuModelName;

    /**
     * 本地大模型-简单流式响应（浏览器输出）
     */
    @RequestMapping("/ollama")
    public Flux<String> ollama(@RequestParam(value = "message") String message, HttpServletResponse response) {
        log.info("message=[{}]", message); // 写一篇关于春天的500字的小作文
        response.setContentType(MediaType.TEXT_EVENT_STREAM_VALUE + "; charset=utf-8");
        return Flux.create(fluxSink -> streamingChatModel2.chat(message, new StreamingChatResponseHandler() {
            @Override
            public void onPartialResponse(String partialResponse) {
                // 每次响应的数据
                //log.info("***AI对话 onNext:\n {}", partialResponse);
                fluxSink.next(partialResponse);
            }

            @Override
            public void onCompleteResponse(ChatResponse completeResponse) {
                // 响应结束的回调
                log.info("***AI对话 onComplete:\n {}", completeResponse.aiMessage().text());
                fluxSink.complete();
            }

            @Override
            public void onError(Throwable error) {
                // 响应出错的回调
                log.error("***AI对话，发送异常", error);
                fluxSink.error(error);
            }
        }));
    }

    /**
     * DeepSeek-简单流式响应（浏览器输出）
     */
    @RequestMapping("/deepseek")
    public Flux<String> deepseek(@RequestParam(value = "message") String message, HttpServletResponse response) {
        log.info("message=[{}]", message); // 写一篇关于春天的200字的小作文
        response.setContentType(MediaType.TEXT_EVENT_STREAM_VALUE + "; charset=utf-8");
        return Flux.create(fluxSink -> streamingChatModel3.chat(message, new StreamingChatResponseHandler() {
            @Override
            public void onPartialResponse(String partialResponse) {
                // 每次响应的数据
                //log.info("***AI对话 onNext:\n {}", partialResponse);
                fluxSink.next(partialResponse);
            }

            @Override
            public void onCompleteResponse(ChatResponse completeResponse) {
                // 响应结束的回调
                log.info("***AI对话 onComplete:\n {}", completeResponse.aiMessage().text());
                fluxSink.complete();
            }

            @Override
            public void onError(Throwable error) {
                // 响应出错的回调
                log.error("***AI对话，发送异常", error);
                fluxSink.error(error);
            }
        }));
    }

    /**
     * 智普AI-简单流式响应（浏览器输出）
     */
    @RequestMapping("/zhipu")
    public Flux<String> zhipu(@RequestParam(value = "message") String message, HttpServletResponse response) {
        log.info("message=[{}]", message); // 写一篇关于春天的200字的小作文
        response.setContentType(MediaType.TEXT_EVENT_STREAM_VALUE + "; charset=utf-8");
        ZhipuAiStreamingChatModel zhipuAiStreamingChatModel = ZhipuAiStreamingChatModel.builder()
                .apiKey(zhipuApiKey)
                .model(zhipuModelName)
                .callTimeout(Duration.ZERO)
                .connectTimeout(Duration.ZERO)
                .writeTimeout(Duration.ZERO)
                .readTimeout(Duration.ZERO)
                .build();
        return Flux.create(fluxSink -> zhipuAiStreamingChatModel.chat(message, new StreamingChatResponseHandler() {
            @Override
            public void onPartialResponse(String partialResponse) {
                // 每次响应的数据
                //log.info("***AI对话 onNext:\n {}", partialResponse);
                fluxSink.next(partialResponse);
            }

            @Override
            public void onCompleteResponse(ChatResponse completeResponse) {
                // 响应结束的回调
                log.info("***AI对话 onComplete:\n {}", completeResponse.aiMessage().text());
                fluxSink.complete();
            }

            @Override
            public void onError(Throwable error) {
                // 响应出错的回调
                log.error("***AI对话，发送异常", error);
                fluxSink.error(error);
            }
        }));
    }

    /**
     * 简单流式响应（浏览器输出）
     */
    @RequestMapping("/chat2")
    public Flux<String> chat2(@RequestParam(value = "message") String message, HttpServletResponse response) {
        log.info("message=[{}]", message); // 写一篇200字的小作文
        response.setContentType(MediaType.TEXT_EVENT_STREAM_VALUE + "; charset=utf-8");
        return Flux.create(fluxSink -> streamingChatModel.chat(message, new StreamingChatResponseHandler() {
            @Override
            public void onPartialResponse(String partialResponse) {
                // 每次响应的数据
                //log.info("***AI对话 onNext:\n {}", partialResponse);
                fluxSink.next(partialResponse);
            }

            @Override
            public void onCompleteResponse(ChatResponse completeResponse) {
                // 响应结束的回调
                log.info("***AI对话 onComplete:\n {}", completeResponse.aiMessage().text());
                // 文章题目：秋天的公园
                // 文章内容：清晨，我漫步在小区附近的公园里。一阵微凉的秋风拂过，金黄的银杏叶纷纷飘落，像一只只翩翩起舞的蝴蝶。草坪上点缀着零星的野花，露珠在阳光下闪烁着晶莹的光芒。远处传来孩子们欢快的笑声，他们在落叶堆里嬉戏玩耍。老人们在小径上散步，或在亭子里打太极。湖面波光粼粼，几只白鹭优雅地掠过水面，留下一道道涟漪。空气中弥漫着淡淡的桂花香，让人心旷神怡。这个季节总让人感到宁静祥和，仿佛整个世界都放慢了脚步。我深深地吸了一口气，感受着这份独属于秋天的美好。
                fluxSink.complete();
            }

            @Override
            public void onError(Throwable error) {
                // 响应出错的回调
                log.error("***AI对话，发送异常", error);
                fluxSink.error(error);
            }
        }));
    }

    /**
     * 简单流式响应（控制台输出）
     */
    @RequestMapping("/chat")
    public String chat(@RequestParam(value = "message") String message) throws InterruptedException {
        log.info("message=[{}]", message); // 写一篇200字的小作文
        CountDownLatch latch = new CountDownLatch(1);
        streamingChatModel.chat(message, new StreamingChatResponseHandler() {
            @Override
            public void onPartialResponse(String partialResponse) {
                // 每次响应的数据
                log.info("***AI对话 onNext:\n {}", partialResponse);
            }

            @Override
            public void onCompleteResponse(ChatResponse completeResponse) {
                // 响应结束的回调
                log.info("***AI对话 onComplete:\n {}", completeResponse.aiMessage().text());
                // 文章题目：秋天的公园
                // 文章内容：清晨，我漫步在城市中心的小公园里。金黄的银杏叶铺满小径，踩上去发出沙沙的声响。秋风轻拂，几片落叶随风起舞，像蝴蝶在空中翩跹。老人们在凉亭里打太极，动作缓慢而有力；孩子们在草坪上追逐嬉戏，欢声笑语回荡在空中。远处的湖面波光粼粼，几只白鹭优雅地掠过水面。这个季节的公园，既有夏日的生机，又添了几分成熟的韵味。在这里，时间仿佛放慢了脚步，让人感受到生活的美好与宁静。我深吸一口气，清新的空气中夹杂着桂花的甜香，令人心旷神怡。这就是秋天的魅力，平凡却动人。
                latch.countDown();
            }

            @Override
            public void onError(Throwable error) {
                // 响应出错的回调
                log.error("***AI对话，发送异常", error);
                latch.countDown();
            }
        });
        latch.await(60, TimeUnit.SECONDS);
        return "流式调用成功...";
    }

}
